import os.path

import numpy as np
import cv2
import imageConversion
import matplotlib.pyplot as plt
from PIL import Image


def clean_nan(image, nan_num):
    """
    根据nan，提取周围部分，并将剩余的nan数据转为nan_num
    Auth: WZW
    Args:
        image:
        nan_num:
    Returns:
        compressed data
    """
    flag = True
    line = 0
    row, col = image.shape
    # 删除全为nan的行
    for i in image:
        if len(i[np.isnan(i)]) < col:
            if flag:
                b = np.array(i)
                flag = False
            else:
                b = np.c_[b, i]
            line = line + 1
    flag = True
    line0 = 0
    # 删除全为列的行
    for i in b:
        if len(i[np.isnan(i)]) < line:
            if flag:
                c = np.array(i)
                flag = False
            else:
                c = np.c_[c, i]
            line0 = line0 + 1
    if not np.isnan(nan_num):
        print("数据压缩，修改nan为：" + str(nan_num))
        nan_index = np.isnan(c)
        c[nan_index] = nan_num
    return c


def clean_light(img_name, targetMachine, img_show=False):
    """
    亮度图清洗，去除黑色部分，只保留工件
    必须传入包括文件名的文件路径
    Args:
        img_name: image name with file path
        imshow:
        imsave:
        target_folder: folder to save image

    Returns:
        rgb image
    """
    image = cv2.imread(img_name)
    if img_name.count("_side_") > 0:
        image = image[1020:1601, 800:1253, :]
    else:
        if targetMachine == 6:
            image = image[1020:1601, 809:1275, :]
        elif targetMachine == 4:
            image = image[1024:1601, 743:1209, :]
        else:
            raise ValueError("工件当前仅支持M4与M6，请输入4/6")
    # cv2.imshow('meta', image)
    # cv2.waitKey(0)
    """
    cut white area
    """
    # if image.shape[0] > 100 and image.shape[1] > 405 and image[100][401][0] == 255 and image[100][401][1] == 255 and \
    #         image[100][401][2] == 255:
    #     image = image[:, :400, :]
    """
    clean pic
    """
    img_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    flag = True
    line = 0
    row, col, chan = image.shape
    # print(image.shape)
    ind = 0
    sumArr = np.zeros(row)
    for i in img_gray:
        sumArr[ind] = np.sum(i)
        if np.sum(i):
            if flag:
                b = np.array(i)
                flag = False
            else:
                b = np.c_[b, i]
            line = line + 1
        ind = ind + 1
    zeroInd = np.where(sumArr != 0)
    img_b = image[zeroInd[0], :, :]

    # print(img_b.shape)
    # print(b.shape)
    # cv2.imshow('img_cut', b)
    # cv2.imshow('rgb', img_b)
    # cv2.waitKey(0)

    row, col, chan = img_b.shape
    del sumArr
    del zeroInd
    flag = True
    line0 = 0
    sumArr = np.zeros(col)
    ind = 0
    for i in b:
        sumArr[ind] = np.sum(i)
        if np.sum(i):
            if flag:
                c = np.array(i)
                flag = False
            else:
                c = np.c_[c, i]
            line0 = line0 + 1
        ind = ind + 1
    zeroInd = np.where(sumArr != 0)
    img_c = img_b[:, zeroInd[0], :]
    # print(img_c.shape)

    if img_show:
        # cv2.imshow('img_cut_c', c)
        # cv2.imshow('rgb_c', img_c)
        # cv2.waitKey(0)
        plt.imshow(img_c)
        plt.show()
    return img_c


def image_narrow_8bit(img_light, img_height):
    """

    Args:
        img_light: nd.array, light image
        img_height: nd.array, height image
    Returns:

    """
    print(img_light.shape)
    print(img_height.shape)
    a, b = img_light.shape[:2]
    img_light = cv2.resize(img_height, (b, a))
    return img_light


def image_blend_alpha_8bit(img_rgb, img_height, alpha=0.5, img_show=False, img_save=False, target_folder="",
                           file_name="out"):
    """
    Auth: WZW
    Args:
        img_rgb: nd.array, light image
        img_height: nd.array, height image
        alpha: 0 - 1, used to blend picture

    Returns:

    """
    print(img_rgb.shape)
    print(img_height.shape)

    img_rgb_pil = Image.fromarray(img_rgb, mode='RGB')
    img_height_pil = Image.fromarray(img_height, mode='RGB')

    img_rgb_alpha = img_rgb_pil.convert('RGBA')
    img_height_alpha = img_height_pil.convert('RGBA')

    blend_img = Image.blend(img_height_alpha, img_rgb_alpha, alpha)
    # img_height_alpha.show()
    if img_show:
        img_rgb_alpha.show()
        blend_img.show()
    if img_save:
        img_rgb_alpha.save(os.path.join(target_folder, file_name + "_light.png"))
        blend_img.save(os.path.join(target_folder, file_name + "_fusion.png"))
    return blend_img


def image_channels_add_8bit(img_rgb, img_height, img_show=False, img_save=False, target_folder="./",
                            file_name="out"):
    """

    Args:
        img_rgb:
        img_height:
        img_show:
        img_save:
        target_folder:
        file_name:

    Returns:

    """
    meta_img = np.zeros(img_rgb.shape)
    img_height_gray = cv2.cvtColor(img_height, cv2.COLOR_RGB2GRAY)
    img_rgb_gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
    meta_img[:, :, 0] = img_height_gray
    meta_img[:, :, 1] = img_rgb_gray
    # meta_img[:, :, 2] = meta_img[:, :, 2] * 50
    meta_img = meta_img.astype(np.uint8)
    if img_show:
        # cv2.imshow("test", meta_img)
        # cv2.waitKey(0)
        plt.imshow(meta_img)
        plt.show()
    if img_save:
        meta_img = cv2.cvtColor(meta_img, cv2.COLOR_BGR2RGB)
        cv2.imwrite(os.path.join(target_folder, file_name + "_channelsBlend.png"), meta_img)
    return


def generate_dataset_image(filePath, imgPath, machineNumber, img_show=False, img_save=False, targetFolder="./"):
    img_light = clean_light(imgPath, machineNumber, img_show=img_show)
    npData = np.load(filePath)
    img_height = clean_nan(npData, np.nan)
    img16, img8, img_max, img_min = imageConversion.outputNorm(img_height, 0)
    img8_rgb = imageConversion.gray2rgb_8bit(img8)
    reshaped_img_height = image_narrow_8bit(img_light, img8_rgb)
    image_blend_alpha_8bit(img_light, reshaped_img_height, img_show=img_show, img_save=img_save,
                           target_folder=targetFolder, file_name=os.path.basename(filePath).split('.')[0])
    image_channels_add_8bit(img_light, reshaped_img_height, img_show=img_show, img_save=img_save,
                            target_folder=targetFolder, file_name=os.path.basename(filePath).split('.')[0])


if __name__ == "__main__":
    """
    单独图片生成
    """
    # imgRootCorrect = r"D:/Programs/data/data/correct/"
    # imgRootError = r"D:/Programs/data/data/err/"
    # # filePath = "D:/Programs/data/data/correct/35051772613505277134A28022200109563_side_2022-02-28-11-22-25.npy"
    # filePath = "D:/Programs/data/data/correct/35051772613505277134A28022200109567_up_2022-02-28-11-30-55.npy"
    # # imgPath = "D:/Programs/data/data/correct/35051772613505277134A28022200109563_side_2022-02-28-11-22-25.png"
    # imgPath = "D:/Programs/data/data/correct/35051772613505277134A28022200109567_up_2022-02-28-11-30-55.png"
    #
    # img_light = clean_light(imgPath, 6)
    # # # image up 201, 446
    # # # image side 398, 101
    #
    # npData = np.load(filePath)
    # img_height = clean_nan(npData, np.nan)
    #
    # # npy up 1391, 3081
    # # npy side 2755, 698
    # # npy / image, up 6.920398009950249, 6.908071748878924
    # # npy / image, side 6.922110552763819, 6.910891089108911
    #
    # img16, img8, img_max, img_min = imageConversion.outputNorm(img_height, 0)
    # img8_rgb = imageConversion.gray2rgb_8bit(img8)
    # reshaped_img_height = image_narrow_8bit(img_light, img8_rgb)
    # # image_blend_alpha_8bit(img_light, reshaped_img_height)
    """
    save fusion image in folder
    """
    # files = os.listdir(imgRootCorrect)
    # for i in files:
    #     if i[-3:] == 'npy':
    #         print(i)
    #         img_light = clean_light(os.path.join(imgRootCorrect, i[:-3] + "png"), imshow=True)
    #         npData = np.load(os.path.join(imgRootCorrect, i))
    #         img_height = clean_nan(npData, np.nan)
    #         img16, img8, img_max, img_min = imageConversion.outputNorm(img_height, 0)
    #         img8_rgb = imageConversion.gray2rgb_8bit(img8, generate_img=True,
    #                                                  target_folder="D:/Programs/data/fusionImage/correct",
    #                                                  fileName=i[:-4] + "_normalRBG.png")
    #         reshaped_img_height = image_narrow_8bit(img_light, img8_rgb)
    #         image_blend_alpha_8bit(img_light, reshaped_img_height, img_save=True,
    #                                target_folder="D:/Programs/data/fusionImage/correct", file_name=i[:-4])

    # files = os.listdir(imgRootError)
    # for i in files:
    #     if i[-3:] == 'npy':
    #         print(i)
    #         img_light = clean_light(os.path.join(imgRootError, i[:-3] + "png"), imshow=True)
    #         npData = np.load(os.path.join(imgRootError, i))
    #         img_height = clean_nan(npData, np.nan)
    #         img16, img8, img_max, img_min = imageConversion.outputNorm(img_height, 0)
    #         img8_rgb = imageConversion.gray2rgb_8bit(img8, generate_img=True,
    #                                                  target_folder="D:/Programs/data/fusionImage/error",
    #                                                  fileName=i[:-4] + "_normalRBG.png")
    #         reshaped_img_height = image_narrow_8bit(img_light, img8_rgb)
    #         image_blend_alpha_8bit(img_light, reshaped_img_height, img_save=True,
    #                                target_folder="D:/Programs/data/fusionImage/error", file_name=i[:-4])

    """
    new channel add
    """
    # image_channels_add_8bit(img_light, reshaped_img_height, img_show=True, img_save=True)
    """
    use generate_dataset_image
    """
    # # img = "D:/Programs/data/dataset20220623/meta/correct/3505277134ML085B085301110012200494_up_2022-02-18-13-10-35.png"
    # img = "D:/Programs/data/dataset20220623/meta/correct/35051772623505277134A09062200143439_up_2022-06-09-12-21-23.png"
    # # npy = "D:/Programs/data/dataset20220623/meta/correct/3505277134ML085B085301110012200494_up_2022-02-18-13-10-35.npy"
    # npy = "D:/Programs/data/dataset20220623/meta/correct/35051772623505277134A09062200143439_up_2022-06-09-12-21-23.npy"
    # generate_dataset_image(npy, img, 6, img_show=True)

    """
    generate dataset
    """
    rootErr = "D:/Programs/data/dataset20220623/meta/error"
    rootCorrect = "D:/Programs/data/dataset20220623/meta/correct"
    rootCorrect01 = "D:/Programs/data/dataset20220623/meta/correct01"
    # files = os.listdir(rootErr)
    # for i in files:
    #     if (i[-3:] == "npy"):
    #         generate_dataset_image(os.path.join(rootErr, i), os.path.join(rootErr, i[:-3] + "png"), 4, img_save=True,
    #                                targetFolder="D:/Programs/data/dataset20220623/generate/error")
    files = os.listdir(rootCorrect)
    for i in files:
        if (i[-3:] == "npy"):
            generate_dataset_image(os.path.join(rootCorrect, i), os.path.join(rootCorrect, i[:-3] + "png"), 4,
                                   img_save=True, targetFolder="D:/Programs/data/dataset20220623/generate/correct")

    files = os.listdir(rootCorrect01)
    for i in files:
        if (i[-3:] == "npy"):
            generate_dataset_image(os.path.join(rootCorrect01, i), os.path.join(rootCorrect01, i[:-3] + "png"), 6,
                                   img_save=True, targetFolder="D:/Programs/data/dataset20220623/generate/correct")
